import numpy as np
from assignment2.python import helper



def distance_cal(im1, im2, xe, ye, pts1):
    # return np.linalg.norm(im1[y - 1:y + 2][x - 1:x + 2], im1[y - 1:y + 2][x - 1:x + 2]).sum()
    xl, yl = int(pts1[0]), int(pts1[1])
    distance_sum = 0
    # print(im2[500][200][2])
    for delta_x in range(-2, 3):
        for delta_y in range(-2, 3):
            distance_sum += np.sqrt((im1[yl + delta_y][xl + delta_x][0] - im2[ye + delta_y][xe + delta_x][0]) ** 2 + (
                    im1[yl + delta_y][xl + delta_x][1] - im2[ye + delta_y][xe + delta_x][1]) ** 2 + (
                                            im1[yl + delta_y][xl + delta_x][2] - im2[ye + delta_y][xe + delta_x][
                                        2]) ** 2)
    return distance_sum


def cal_one(im1, im2, F, pts1, size):
    line = F @ pts1
    mini = float(np.inf)
    pair = ()
    if line[1] != 0:
        for x in range(2, size[1] - 2):
            y = int(np.floor(-(line[0] * x + line[2]) / line[1]))
            if y == 0 or y == size[0]-1:
                continue
            val = distance_cal(im1, im2, x, y, pts1)
            if mini > val:
                mini = val
                pair = (y, x)
    else:
        x = int(round(- line[2] / line[0]))
        for y in range(2, size[1] - 2):
            val = distance_cal(im1, im2, x, y, pts1)
            if mini > val:
                mini = val
                pair = (y, x)
    return np.array([[pair[1], pair[0]]], dtype=np.float32)


def epipolar_correspondences(im1, im2, F, pts1):
    size = im1.shape
    shape = pts1.shape
    # print(pts1)
    # add '1' column and get 3*N matrix for multiply
    add = np.ones((1, pts1.shape[0]), dtype=np.float32, order='C')
    # print(shape)
    pts1 = np.insert(pts1, 2, values=add, axis=1).T
    pts2 = np.empty(shape, dtype=float, order='C')
    # print(pts1[:,2])
    for t in range(0, shape[0]):
        one_pair = cal_one(im1, im2, F, pts1[:, t], size)
        pts2[t] = one_pair

    return pts2
